import numpy as np
import torch
import os
from tqdm import tqdm
from torch import nn
from torch import optim
from config import Config
from preprocess import preprocess
from TextCNN import TextCNN
import torchnet.meter as meter
from torchsummary import summary
import matplotlib.pyplot as plt


def train(model, train_loader, dev_loader, test_loader):
    model.train()

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=Config.learning_rate)
    loss_meter = meter.AverageValueMeter()
    acc_meter = meter.ClassErrorMeter(accuracy=True)
    loss_list, loss_val_list, acc_list, acc_val_list = [], [], [], []
    best_validation_acc = 0
    if os.path.exists(Config.save_dir_loss):
        loss_list = torch.load(Config.save_dir_loss)
    if os.path.exists(Config.save_dir_loss_val):
        loss_val_list = torch.load(Config.save_dir_loss_val)
    if os.path.exists(Config.save_dir_acc):
        acc_list = torch.load(Config.save_dir_acc)
    if os.path.exists(Config.save_dir_acc_val):
        acc_val_list = torch.load(Config.save_dir_acc_val)

    for epoch in range(Config.epoch):
        loss_meter.reset()
        acc_meter.reset()
        loop = tqdm(enumerate(train_loader), total=len(train_loader))
        for i, batch_data in loop:
            input_data, label = batch_data[0].to(device), batch_data[1].to(device)
            optimizer.zero_grad()
            outputs = model(input_data)
            loss = criterion(outputs, label)
            loss.backward()
            optimizer.step()
            loss_meter.add(loss.item())
            acc_meter.add(outputs.cpu().detach().numpy(), label)
            # set progress bar
            loop.set_description(f'TRAIN Epoch [{epoch}/{Config.epoch}]')
            loop.set_postfix(loss=loss_meter.mean, acc=acc_meter.value()[0])
        loss_list.append(loss_meter.mean)
        acc_list.append(acc_meter.value())
        model.save(Config.save_dir_param)

        val_loss, val_acc = validate(model, dev_loader, criterion)
        if val_acc > best_validation_acc:
            best_validation_acc = val_acc
            test_acc = test(model, test_loader)
            if test_acc > 0.83:
                model.save('param_{}.pth'.format(epoch))
        model.train()
        loss_val_list.append(val_loss)
        acc_val_list.append(val_acc)
    torch.save(loss_list, Config.save_dir_loss)
    torch.save(acc_list, Config.save_dir_acc)
    torch.save(loss_val_list, Config.save_dir_loss_val)
    torch.save(acc_val_list, Config.save_dir_acc_val)


def validate(model, dev_loader, criterion):
    model.eval()
    loss_meter = meter.AverageValueMeter()
    acc_meter = meter.ClassErrorMeter(accuracy=True)
    loop = tqdm(enumerate(dev_loader), total=len(dev_loader))
    for i, batch_data in loop:
        input_data, label = batch_data[0].to(device), batch_data[1].to(device)
        outputs = model(input_data)
        loss = criterion(outputs, label)
        loss_meter.add(loss.item())
        acc_meter.add(outputs.cpu().detach().numpy(), label)
        # set progress bar
        loop.set_description(f'VALIDATION ')
        loop.set_postfix(loss=loss_meter.mean, acc=acc_meter.value()[0])
    return loss_meter.mean, acc_meter.value()[0]


def test(model, test_loader):
    model.eval()
    criterion = nn.CrossEntropyLoss()
    loss_meter = meter.AverageValueMeter()
    TP, TN, FP, FN = 0, 0, 0, 0
    loop = tqdm(enumerate(test_loader), total=len(test_loader))
    for i, batch_data in loop:
        input_data, label = batch_data[0].to(device), batch_data[1].to(device)
        outputs = model(input_data)
        loss = criterion(outputs, label)
        loss_meter.add(loss.item())
        pred = outputs.argmax(axis=1).cpu()
        TP += np.logical_and(pred == 0, label.cpu() == 0).sum()
        TN += np.logical_and(pred == 1, label.cpu() == 1).sum()
        FP += np.logical_and(pred == 0, label.cpu() == 1).sum()
        FN += np.logical_and(pred == 1, label.cpu() == 0).sum()
        # set progress bar
        loop.set_description(f'TEST ')
        loop.set_postfix(loss=loss_meter.mean)
    Precision = TP / (TP + FP + 1e-6)
    Recall = TP / (TP + FN + 1e-6)
    F1 = 2 * Precision * Recall / (Precision + Recall + 1e-6)
    Accuracy = (TP + TN) / (TP + TN + FP + FN + 1e-6)
    print('TEST RESULT\n'
          '-----------\n'
          'Accuracy: {}\n'
          'Precision: {}\n'
          'Recall: {}\n'
          'F1 Score: {}\n'
          'loss: {}\n'
          'Confusion Matrix:\n'
          '|{} {}|\n'
          '|{} {}|'.format(Accuracy, Precision, Recall, F1, loss_meter.mean, TP, FP, FN, TN))
    return Accuracy


def draw_curves():
    def fix(x):
        if type(x) == list:
            x = x[0]
        return x

    tr_loss = torch.load(Config.save_dir_loss)
    tr_acc = [x[0] for x in torch.load(Config.save_dir_acc)]
    dev_loss = torch.load(Config.save_dir_loss_val)
    dev_acc = [fix(x) for x in torch.load(Config.save_dir_acc_val)]
    X = torch.arange(1, len(tr_loss) + 1)
    print(tr_acc)
    print(tr_loss)
    print(dev_loss)
    print(dev_acc)

    # draw loss curves
    plt.plot(X, tr_loss, label='train loss')
    plt.plot(X, dev_loss, label='validation loss')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend()
    plt.title('Train/Validation Loss')
    plt.show()

    # draw acc curves
    plt.plot(X, tr_acc, label='train accuracy')
    plt.plot(X, dev_acc, label='validation accuracy')
    plt.xlabel('epoch')
    plt.ylabel('accuracy')
    plt.legend()
    plt.title('Train/Validation Accuracy')
    plt.show()


def main():
    train_loader, dev_loader, test_loader = preprocess()
    model = TextCNN(
        Config.vocab_size,
        Config.embedding_dim,
        Config.num_filters,
        Config.kernel_size,
        Config.dropout_keep_prob
    ).to(device)
    if os.path.exists(Config.save_dir_param):
        model.load_state_dict(torch.load(Config.save_dir_param))

    # train(model, train_loader, dev_loader, test_loader)
    if os.path.exists(Config.save_dir_param):
        model.load_state_dict(torch.load('param_9.pth'))
    test(model, test_loader)
    # print(summary(model, (75,), batch_size=128, device=device))
    # draw_curves()


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    main()
